#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @author: Gaoxiang

import torch
from sklearn.metrics import classification_report
from text_cls_dataloader import DataIterator


class Evaluator:
    def __init__(self, config, model, logger, datafield):
        self.config = config
        self.model = model
        self.logger = logger
        self.valid_data = DataIterator(raw_set=config["valid_data_path"],
                                       is_test_set=False,
                                       use_cuda=config["use_cuda"],
                                       batchsize=config["batch_size"],
                                       datafield=datafield,
                                       config=config)
        self.class_num = config["num_class"]
        self.stats_dict = {"total_correct": 0, "total_wrong": 0}  # 用于存储测试结果
        for i in range(self.class_num):
            self.stats_dict['{}_correct'.format(i)] = 0
            self.stats_dict['{}_wrong'.format(i)] = 0

    def eval(self, epoch, config):
        self.logger.info("开始测试第%d轮模型效果：" % epoch)
        self.model.eval()
        self.stats_dict = {"total_correct": 0, "total_wrong": 0}  # 用于存储测试结果
        for i in range(self.class_num):
            self.stats_dict['{}_correct'.format(i)] = 0
            self.stats_dict['{}_wrong'.format(i)] = 0
        for index, data in enumerate(self.valid_data.dataset):
            if config["use_cuda"]:
                content, label = data.text.cuda(), data.label.cuda()
            else:
                content, label = data.text, data.label

            with torch.no_grad():
                pred_results = self.model(content)  # 不输入labels，使用模型当前参数进行预测
            self.write_stats(label, pred_results)
        acc = self.show_stats()
        return acc

    def eval_sklearn(self, epoch, config):
        self.logger.info("开始测试第%d轮模型效果：" % epoch)
        self.model.eval()

        for index, data in enumerate(self.valid_data.get_tier_data()):
            if config["use_cuda"]:
                content, label = data.text.cuda(), data.label.cuda()
            else:
                content, label = data.text, data.label

            with torch.no_grad():
                pred_results = self.model(content)
            pred_results = torch.argmax(input=pred_results, dim=1)

            if index == 0:
                label_total = label
                pred_total = pred_results
            else:
                label_total = torch.cat((label_total, label))
                pred_total = torch.cat((pred_total, pred_results))
        self.logger.info("\n" + classification_report(label_total, pred_total))

    def write_stats(self, labels, pred_results):
        assert len(labels) == len(pred_results)
        for true_label, pred_label in zip(labels, pred_results):
            pred_label = torch.argmax(pred_label)
            if int(true_label) == int(pred_label):
                self.stats_dict["total_correct"] += 1
                self.stats_dict['{}_correct'.format(true_label.item())] += 1
            else:
                self.stats_dict["total_wrong"] += 1
                self.stats_dict['{}_wrong'.format(true_label.item())] += 1
        return

    def show_stats(self):
        correct = self.stats_dict["total_correct"]
        wrong = self.stats_dict["total_wrong"]
        self.logger.info("预测集合条目总量：%d" % (correct +wrong))
        self.logger.info("预测正确条目：%d，预测错误条目：%d" % (correct, wrong))
        self.logger.info("预测准确率：%f" % (correct / (correct + wrong)))
        for i in range(self.class_num):
            i_correct = self.stats_dict['{}_correct'.format(i)]
            i_wrong = self.stats_dict['{}_wrong'.format(i)]
            self.logger.info("{}_预测正确条目：%d，{}_预测错误条目：%d".format(i, i) % (i_correct, i_wrong))
            if (i_correct + i_wrong) != 0:
                self.logger.info("{}_预测精确度：%f".format(i) % (i_correct / (i_correct + i_wrong)))
        self.logger.info("--------------------")
        return correct / (correct + wrong)